from pathlib import Path
from collections import defaultdict

dataset_path_dict = {
    'car': '/home/zgp/Datasets/cars196_notbyclass',
    'cub': '/home/zgp/Datasets/cub200',
    'shop': '/home/zgp/Datasets/in-shop',
    'sop': '/home/zgp/Datasets/online_products'
}

def gen_cub_dict():
    rv = {'tra': defaultdict(list), 'test': defaultdict(list)}
    dataset_path = Path(dataset_path_dict['cub'])
    trainlist = dataset_path.joinpath('train100.txt')
    testlist = dataset_path.joinpath('test100.txt')
    with open(trainlist, 'r') as f:
        for l in f:
            id, relative_path = l.split()
            class_id = relative_path.split('.')[0]
            key = f'tra_{class_id}'
            rv['tra'][key].append(str(dataset_path.joinpath('images', relative_path)))

    with open(testlist, 'r') as f:
        for l in f:
            id, relative_path = l.split()
            class_id = relative_path.split('.')[0]
            key = f'test_{class_id}'
            rv['test'][key].append(str(dataset_path.joinpath('images', relative_path)))

    return rv

def gen_car_dict():
    rv = {'tra': defaultdict(list), 'test': defaultdict(list)}
    dataset_path = Path(dataset_path_dict['car'])
    # img_path = dataset_path.joinpath('car_ims')
    trainlist = dataset_path.joinpath('train.txt')
    testlist = dataset_path.joinpath('test.txt')
    with open(trainlist, 'r') as f:
        for l in f:
            relative_path, class_id = l.split()
            key = f'tra_{class_id}'
            rv['tra'][key].append(str(dataset_path.joinpath(relative_path)))

    with open(testlist, 'r') as f:
        for l in f:
            relative_path, class_id = l.split()
            key = f'test_{class_id}'
            rv['test'][key].append(str(dataset_path.joinpath(relative_path)))

    return rv

def gen_sop_dict():
    rv = {'tra': defaultdict(list), 'test': defaultdict(list)}
    dataset_path = Path(dataset_path_dict['sop'])
    trainlist = dataset_path.joinpath('Info_Files', 'train.txt')
    testlist = dataset_path.joinpath('Info_Files', 'test.txt')
    with open(trainlist, 'r') as f:
        for l in f:
            relative_path, id = l.split()
            key = f'tra_{id}'
            rv['tra'][key].append(str(dataset_path.joinpath('images', relative_path)))

    with open(testlist, 'r') as f:
        for l in f:
            relative_path, id = l.split()
            key = f'test_{id}'
            rv['test'][key].append(str(dataset_path.joinpath('images', relative_path)))

    return rv

# TODO: gallery.txt
def gen_shop_dict():
    rv = {'tra': defaultdict(list), 'test': defaultdict(list)}
    dataset_path = Path(dataset_path_dict['shop'])
    trainlist = dataset_path.joinpath('Eval', 'train.txt')
    testlist = dataset_path.joinpath('Eval', 'query.txt')
    with open(trainlist, 'r') as f:
        for l in f:
            relative_path, class_id = l.split()
            key = f'tra_{class_id}'
            rv['tra'][key].append(str(dataset_path.joinpath(relative_path)))

    with open(testlist, 'r') as f:
        for l in f:
            relative_path, class_id = l.split()
            key = f'test_{class_id}'
            rv['test'][key].append(str(dataset_path.joinpath(relative_path)))

    return rv

def gen_datadict(name='CAR'):
    if name == 'CAR':
        return gen_car_dict()
    elif name == 'CUB':
        return gen_cub_dict()
    elif name == 'SOP':
        return gen_sop_dict()
    elif name == 'ISC':
        return gen_shop_dict()


if __name__ == "__main__":
    datadict = gen_datadict('ISC')
    # print(datadict)
    assert 'tra' in datadict and 'test' in datadict
    print(datadict['tra'].keys())
    print(datadict['test'].keys())
    __import__('ipdb').set_trace()

    datadict = gen_datadict()
    # print(datadict)
    assert 'tra' in datadict and 'test' in datadict
    print(datadict['tra'].keys())
    print(datadict['test'].keys())

    datadict = gen_datadict(name='CUB')
    assert 'tra' in datadict and 'test' in datadict
    print(datadict['tra'].keys())
    print(datadict['test'].keys())

    datadict = gen_datadict(name='SOP')
    assert 'tra' in datadict and 'test' in datadict
    print(datadict['tra'].keys())
    print(datadict['test'].keys())
